from event_handle.generate import generate

from event_handle.query_rewrite import query_rewrite
from event_handle.recall import recall
from event_handle.rerank import rerank


class EventHandlePipeline:
    def __init__(self, query):
        self.query = query

    async def process(self, stream=False):
        rewrite_result = await query_rewrite(self.query)
        recall_result = recall(rewrite_result)
        rerank_result = rerank(recall_result)
        if not stream:
            return await generate(rerank_result)
        else:
            return await generate(rerank_result, stream=True)
